#pragma once

#include "ATen/core/TensorBody.h"
#include "torch/nn/modules/linear.h"
#include <torch/torch.h>

class Net : public torch::nn::Module
{
private:
    torch::nn::Linear input_{nullptr};
    torch::nn::Linear hidden_{nullptr};
    torch::nn::Linear output_{nullptr};

public:
    Net();

    ~Net();

    at::Tensor forward(at::Tensor x);
};

class CustomDtaset : public torch::data::Dataset<CustomDtaset>
{
private:
    torch::Tensor inputs_; ///< 输入
    torch::Tensor labels_; ///< 标签

public:
    CustomDtaset(torch::Tensor input_data, torch::Tensor label_data);
    torch::data::Example<> get(size_t index) override;
    torch::optional<size_t> size() const override;
};
